import random
from threading import local
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
from isaacgym import gymtorch
from math import sqrt
import math
from omegaconf import read_write
from sympy import false, re
import torch
import cv2
import gym
from gym import spaces

from draw import *

from pcgworker.PCGWorker import *

from stable_baselines3.common.env_checker import check_env

# https://stable-baselines3.readthedocs.io/en/master/guide/custom_env.html

N_DISCRETE_ACTIONS = 6
N_CHANNELS = 3
HEIGHT = 84
WIDTH = 84

class CustomEnv(gym.Env):
    """Custom Environment that follows gym interface"""
    metadata = {'render.modes': ['human', 'rgb_array']}

    def __init__(self, headless_: bool = True, seed_pth: str = "blank_wave.json"):
        super(CustomEnv, self).__init__()
        # Define action and observation space
        self.action_space = spaces.Discrete(N_DISCRETE_ACTIONS)
        self.observation_space = spaces.Box(low=0, high=255, shape=(HEIGHT, WIDTH, N_CHANNELS), dtype=np.uint8)

        torch.manual_seed(3407)

        '''
        simulatiom parameters
        '''
        self.num_matrix_envs = 1
        self.headless = headless_
        self.spd = 10
        max_actors_num_per_env = 9*9*4+1

        self.initial_height = 11

        # get default set of parameters
        sim_params = gymapi.SimParams()
        args = gymutil.parse_arguments()
        sim_params.physx.solver_type = 1
        sim_params.physx.num_position_iterations = 2
        sim_params.physx.num_velocity_iterations = 1
        sim_params.physx.num_threads = args.num_threads
        sim_params.physx.use_gpu = args.use_gpu
        sim_params.use_gpu_pipeline = False #args.use_gpu_pipeline
        sim_params.substeps = 2
        sim_params.dt = 1.0 / 30.0
        # setting up the Z-up axis
        sim_params.up_axis = gymapi.UP_AXIS_Z
        sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)

        # configure the ground plane
        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0, 0, 1) # z-up!
        plane_params.distance = 1
        plane_params.static_friction = 1
        plane_params.dynamic_friction = 1
        plane_params.restitution = 0.1

        # set up the env grid
        num_envs = 1
        spacing = 20
        env_lower = gymapi.Vec3(-spacing, -spacing, -spacing)
        env_upper = gymapi.Vec3(spacing, spacing, spacing)

        # Attractor setup
        attractor_properties = gymapi.AttractorProperties()
        attractor_properties.stiffness = 5e10
        attractor_properties.damping = 0
        attractor_properties.axes = gymapi.AXIS_ROTATION
        attractor_pose = gymapi.Transform()
        attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
        attractor_pose.r = gymapi.Quat(0,1,0,1)

        # create procedural asset
        asset_options = gymapi.AssetOptions()
        asset_options.density = 1.0
        asset_options.fix_base_link = False
        asset_options.linear_damping = 0.1
        asset_options1 = gymapi.AssetOptions()
        asset_options1.density = 0.02
        asset_options1.fix_base_link = True
        asset_options1.linear_damping = 0.1

        color_red = gymapi.Vec3(1,0,0)
        color_green = gymapi.Vec3(0,1,0)

        # load PCG asset
        asset_root = "./assets"
        pt_gray_cube = "PCG/gray_cube.urdf"
        pt_blue_cube = "PCG/blue_cube.urdf"
        pt_red_cube = "PCG/red_cube.urdf"
        pt_white_cube = "PCG/white_cube.urdf"
        pt_yellow_cube = "PCG/yellow_cube.urdf"
        pt_blue_ramp = "PCG/blue_ramp.urdf"
        pt_red_ramp = "PCG/red_ramp.urdf"
        pt_white_ramp = "PCG/white_ramp.urdf"
        pt_yellow_ramp = "PCG/yellow_ramp.urdf"

        # ramp and corner attitutde
        ramp_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0*math.pi)
        ramp_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi)
        ramp_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1*math.pi)
        ramp_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1.5*math.pi)
        corner_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0*math.pi)
        corner_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi)
        corner_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1*math.pi)
        corner_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1.5*math.pi)

        # birick size
        side_l = 2

        # camera sensor properties
        camera_properties = gymapi.CameraProperties()
        camera_properties.width = WIDTH
        camera_properties.height = HEIGHT
        camera_properties.enable_tensors = True

        # rigid shape material properties
        shape_props = [gymapi.RigidShapeProperties()]
        shape_props[0].friction = -1
        shape_props[0].rolling_friction = 1
        shape_props[0].torsion_friction = 1
        shape_props[0].compliance = 0
        shape_props[0].restitution = 0
        shape_props1 = [gymapi.RigidShapeProperties()]
        shape_props1[0].friction = 0.2
        shape_props1[0].rolling_friction = -1
        shape_props1[0].torsion_friction = -1
        shape_props1[0].compliance = 0
        shape_props1[0].restitution = 0
        shape_props2 = [gymapi.RigidShapeProperties()]
        shape_props2[0].friction = 1
        shape_props2[0].rolling_friction = -1
        shape_props2[0].torsion_friction = -1
        shape_props2[0].compliance = 0
        shape_props2[0].restitution = 0

        # single world camera image
        self.camera_sensor_image = np.zeros((HEIGHT,WIDTH, 3), np.uint8)
        # human control
        self.interaction = 0

        # proximity threshold for reward calculation
        self.proximity_threshold = 1.5
        self.max_steps = 1000

        self.show_action = False

        '''
        generate WFC maps
        '''
        PCGWorker_ = PCGWorker(9,9)
        seed = PCGWorker_.from_file(filename = seed_pth)

        # seed_ = PCGWorker_.from_file(filename = "blank_wave.json")
        # seed = PCGWorker_.mutate(seed_,81)
        # PCGWorker_.to_file(seed,filename = "seed.json")
        self.connectivity_map , self.connectivity_map_visual = PCGWorker_.connectivity_analysis(seed,visualize_ = False)

        seeds = [seed]
        for i in range(0,self.num_matrix_envs-1):
            seed_ = PCGWorker_.mutate(seeds[-1],81)
            seeds.append(seed_)
        building_block_pointer = [0,0,0,0,0,0,0,0,0,0,0,0,0]    # 0-4,brick; 5-8,corners; 9-12 ,ramps;

        self.gym = gymapi.acquire_gym()

        self.sim = self.gym.create_sim(args.compute_device_id, args.graphics_device_id, gymapi.SIM_PHYSX, sim_params)

        print("args.graphics_device_id : ",args.graphics_device_id)
        print("args.compute_device_id : ",args.compute_device_id)


        # create the ground plane
        self.gym.add_ground(self.sim, plane_params)

        capsule_asset = self.gym.create_capsule(self.sim, 1, 1, asset_options)
        # capsule_asset = self.gym.create_box(self.sim, 1, 1, 1, asset_options)

        asset_gray_cube = self.gym.load_asset(self.sim, asset_root, pt_gray_cube,asset_options1)
        asset_blue_cube = self.gym.load_asset(self.sim, asset_root, pt_blue_cube,asset_options1)
        asset_red_cube = self.gym.load_asset(self.sim, asset_root, pt_red_cube,asset_options1)
        asset_white_cube = self.gym.load_asset(self.sim, asset_root, pt_white_cube,asset_options1)
        asset_yellow_cube = self.gym.load_asset(self.sim, asset_root, pt_yellow_cube,asset_options1)
        asset_blue_ramp = self.gym.load_asset(self.sim, asset_root, pt_blue_ramp,asset_options1)
        asset_red_ramp = self.gym.load_asset(self.sim, asset_root, pt_red_ramp,asset_options1)
        asset_white_ramp = self.gym.load_asset(self.sim, asset_root, pt_white_ramp,asset_options1)
        asset_yellow_ramp = self.gym.load_asset(self.sim, asset_root, pt_yellow_ramp,asset_options1)

        box_asset = self.gym.create_box(self.sim, 1, 1, 1, asset_options)

        if not self.headless:
            cam_props = gymapi.CameraProperties()
            self.viewer = self.gym.create_viewer(self.sim, cam_props)

        self.envs = []
        self.camera_handles = []
        self.actor_handles = []
        self.food_handles = []
        self.attractor_handles = []
        self.force_sensors = []

        for n in range(self.num_matrix_envs):
            # create environments
            env_ = self.gym.create_env(self.sim, env_lower, env_upper, 4)
            self.envs.append(env_)

            n_actors = 0
            # create actors based on WFC seed
            for i in range(0,9):
                for j in range(0,9):
                    tile_ = seeds[n].wave_oriented[j*9+i][0][0]
                    rot = seeds[n].wave_oriented[j*9+i][0][1]
                    pose_ij = gymapi.Transform()
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 0*side_l)
                    pose_ij.r = gymapi.Quat(0,0,0,1)
                    actor = self.gym.create_actor(env_, asset_gray_cube, pose_ij, "gray_cube" + str(i) + str(j), n, 0)
                    # cubes
                    if tile_<=5:
                        if tile_==1:
                            pass
                        elif tile_==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        elif tile_==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        elif tile_==4:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                        elif tile_==5:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            actor = self.gym.create_actor(env_, asset_white_cube, pose_ij, "white_cube" + str(i) + str(j), n, 0)
                    # corners
                    elif tile_==6:
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    elif tile_==7:
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    elif tile_==8:
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    elif tile_==9:
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                        actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    # ramps
                    elif tile_==10:
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
                    elif tile_==11:
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
                    elif tile_==12:
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
                    elif tile_==13:
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                        actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)

            # generate random x,y coordinates in range [0,18] for the actor
            x = random.uniform(0,16)
            y = random.uniform(0,16)
            facing_ = 0.25 * (float)(random.randint(0,8))
            r_ = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing_*math.pi)

            # create capsule_asset actor in the environment
            pose = gymapi.Transform()
            pose.p = gymapi.Vec3(x, y, self.initial_height)
            pose.r = gymapi.Quat(0,1,0,1)
            cap_handle = self.gym.create_actor(env_, capsule_asset, pose, 'capsule_asset', n, 0)
            self.gym.set_actor_rigid_shape_properties(env_, cap_handle, shape_props2)
            self.gym.set_actor_scale(env_, cap_handle, 0.4)
            self.actor_handles.append(cap_handle)
            self.gym.set_rigid_body_color(env_, cap_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color_red)

            # names = self.gym.get_actor_rigid_body_names(env_, cap_handle)
            # print(names)
            
            # set random position for food
            x = random.uniform(0,16)
            y = random.uniform(0,16)
            # create capsule_asset food in the environment
            pose = gymapi.Transform()
            pose.p = gymapi.Vec3(x, y, self.initial_height)
            pose.r = gymapi.Quat(0,1,0,1)
            food_handle = self.gym.create_actor(env_, capsule_asset, pose, 'food_asset', n, 0)
            self.gym.set_actor_rigid_shape_properties(env_, food_handle, shape_props2)
            self.gym.set_actor_scale(env_, food_handle, 0.9)
            self.food_handles.append(food_handle)
            self.gym.set_rigid_body_color(env_, food_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color_green)


            att_pose = gymapi.Transform()
            att_pose.p = gymapi.Vec3(x,y, self.initial_height)
            att_pose.r = r_
            attractor_properties_ = gymapi.AttractorProperties()
            attractor_properties_.stiffness = 5e10
            attractor_properties_.damping = 0
            attractor_properties_.axes = gymapi.AXIS_ROTATION
            attractor_properties_.target = att_pose
            attractor_properties_.rigid_handle = cap_handle
            attractor_handle_ = self.gym.create_rigid_body_attractor(env_, attractor_properties_)

            attractor_properties_food = gymapi.AttractorProperties()
            attractor_properties_food.stiffness = 5e10
            attractor_properties_food.damping = 0
            attractor_properties_food.axes = gymapi.AXIS_ROTATION
            attractor_properties_food.target = pose
            attractor_properties_food.rigid_handle = food_handle
            attractor_handle_food = self.gym.create_rigid_body_attractor(env_, attractor_properties_food)
            
            h1 = self.gym.create_camera_sensor(env_, camera_properties)
            camera_offset = gymapi.Vec3(-3, 0, 0)
            camera_rotation = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 0, 1), 0.4*math.pi)
            body_handle = self.gym.get_actor_rigid_body_handle(env_, cap_handle, 0)
            self.gym.attach_camera_to_body(h1, env_, body_handle, gymapi.Transform(camera_offset, camera_rotation), gymapi.FOLLOW_TRANSFORM)
            self.camera_handles.append(h1)
            self.attractor_handles.append(attractor_handle_)

            actor_count = self.gym.get_actor_count(env_)
            # # create blank actors
            # num_blk_act = max_actors_num_per_env - actor_count
            # for i in range(num_blk_act):
            #     pose = gymapi.Transform()
            #     pose.p = gymapi.Vec3(0, 0, 0)
            #     pose.r = gymapi.Quat(0,1,0,1)
            #     actor = self.gym.create_actor(env_, box_asset, pose, 'box_asset' + str(i), n, 0)
            #     self.gym.set_actor_scale(env_, actor, 0.000001)

        # a must-have operation 
        self.gym.prepare_sim(self.sim)

        self.step_count = 0

        if not self.headless:
            cam_pos = gymapi.Vec3(0.0, 15.0, 15.0)
            cam_target = gymapi.Vec3(1.5, 4, 0)
            self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

        # time test
        self.start_time = time.time()
        self.last_frame_cnt = 0

        # Create helper geometry used for visualization
        # Create an wireframe axis
        self.axes_geom = gymutil.AxesGeometry(2)

        self.facing = 0
        self.facing_step = 0.25

        # local coordinate system
        self.verts = np.empty((3, 1), gymapi.Vec3.dtype)
        self.verts[0][0] = (1, 0, 0)
        self.verts[1][0] = (0, 1, 0)
        self.verts[2][0] = (0, 0, 1)

        # store initial status for all envs
        self.sim_status = []
        for i in range(self.num_matrix_envs):
            env_status = np.copy(self.gym.get_env_rigid_body_states(self.envs[i], gymapi.STATE_ALL))
            self.sim_status.append(env_status)
        
    def set_world(self, seed_):

        pass

    def calculate_reward(self, env_, actor_handle, food_handle):

        i = 0
        body_states_actor = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states_actor['pose']['p']
        body_states_food = self.gym.get_actor_rigid_body_states(self.envs[i], self.food_handles[i], gymapi.STATE_ALL)
        food_p = body_states_food['pose']['p']

        # calculate distance between actor and food
        ap = np.array([actor_p[0][0], actor_p[0][1], 0])
        fp = np.array([food_p[0][0], food_p[0][1], 0])
        dist = np.linalg.norm(ap - fp)
        # print(dist)

        if dist < self.proximity_threshold:
            # print('reward: 1')
            return 1
        else:
            return 0

    def done_determinant(self, env_, actor_handle, food_handle, reward = 0, step = 0):

        if reward == 1:
            return True

        i = 0
        body_states_actor = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states_actor['pose']['p']
        actor_cnnc_x = ((actor_p[0][0]+1)/(9*2))*20*9
        actor_cnnc_y = ((actor_p[0][1]+1)/(9*2))*20*9

        body_states_food = self.gym.get_actor_rigid_body_states(self.envs[i], self.food_handles[i], gymapi.STATE_ALL)
        food_p = body_states_food['pose']['p']
        food_cnnc_x = ((food_p[0][0]+1)/(9*2))*20*9
        food_cnnc_y = ((food_p[0][1]+1)/(9*2))*20*9

        if not (actor_cnnc_x >= 0 and actor_cnnc_x < 20*9 and actor_cnnc_y >= 0 and actor_cnnc_y < 20*9):
            return True
        
        if not (food_cnnc_x >= 0 and food_cnnc_x < 20*9 and food_cnnc_y >= 0 and food_cnnc_y < 20*9):
            return True

        actor_cnnt_idx = self.connectivity_map[int(actor_cnnc_y)][int(actor_cnnc_x)]
        food_cnnt_idx = self.connectivity_map[int(food_cnnc_y)][int(food_cnnc_x)]

        if actor_cnnt_idx != 1 and food_cnnt_idx != 1:
            if actor_cnnt_idx != food_cnnt_idx:
                return True

        if step > self.max_steps:
            return True

        return False

    def step(self, action : int, physics_only = False):

        self.step_count += 1

        l_ = int((HEIGHT-1)/4)

        i = 0

        '''
        apply action
        '''
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        # transform global coordinate system to local coordinate system
        body_ = self.gym.get_actor_rigid_body_handle(self.envs[i], self.actor_handles[i], 0)
        body_t = self.gym.get_rigid_transform(self.envs[i], body_)
        body_t.p = gymapi.Vec3(0,0,0)
        verts_ = body_t.transform_points(self.verts)

        # maintain vertical speed
        z_spd = body_states['vel']['linear'][0][2]

        if action == 0:
            body_states['vel']['linear'].fill((self.spd * verts_[1][0][0],self.spd * verts_[1][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 1:
            body_states['vel']['linear'].fill((-self.spd * verts_[1][0][0],-self.spd * verts_[1][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 2:
            body_states['vel']['linear'].fill((-self.spd * verts_[2][0][0],-self.spd * verts_[2][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 3:
            body_states['vel']['linear'].fill((self.spd * verts_[2][0][0],self.spd * verts_[2][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 4:
            self.facing -= self.facing_step
            attractor_pose = gymapi.Transform()
            attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
            attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), self.facing*math.pi)
            self.gym.set_attractor_target(self.envs[i], self.attractor_handles[i], attractor_pose)
        elif action == 5:
            self.facing += self.facing_step
            attractor_pose = gymapi.Transform()
            attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
            attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), self.facing*math.pi)
            self.gym.set_attractor_target(self.envs[i], self.attractor_handles[i], attractor_pose)
        # # break !!
        # elif action == 6:
        #     body_states['vel']['linear'].fill((0,0,0))
        #     self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)

        observation = np.zeros((HEIGHT,WIDTH, 4), np.uint8)

        '''
        step the physics
        '''
        self.gym.simulate(self.sim)
        self.gym.fetch_results(self.sim, True)
        if not self.headless:
            self.gym.clear_lines(self.viewer)
        self.gym.step_graphics(self.sim)

        if not physics_only:

            self.gym.render_all_camera_sensors(self.sim)

            '''
            get observation
            '''
            rgb_image = self.gym.get_camera_image(self.sim, self.envs[i], self.camera_handles[i], gymapi.IMAGE_COLOR)
            # create an opencv image of size (HEIGHT,WIDTH , 4)
            
            observation[:, :, :4] = rgb_image.reshape((HEIGHT,WIDTH, 4))
            # rotate 90 degrees counter-clockwise
            observation = np.rot90(observation, -1)
            # transform the image from RGBA to RGB
            observation = observation[:, :, :3]

            # velocity indicator
            body_ti = body_t.inverse()
            vel_ = np.empty((1, 1), gymapi.Vec3.dtype)
            vel_[0][0] = (body_states['vel']['linear'][0][0], body_states['vel']['linear'][0][1], body_states['vel']['linear'][0][2])
            vel_ego = body_ti.transform_points(vel_)
            # print(vel_ego)
            if vel_ego[0][0][1]>=0 :
                h_ = l_*(vel_ego[0][0][1]/self.spd)
                h_ = min(l_,h_)
                observation[:12,:(int)(h_),:]=(200,200,200)
            elif vel_ego[0][0][1]<0 :
                h_ = -l_*(vel_ego[0][0][1]/self.spd)
                h_ = min(l_,h_)
                observation[:12,l_:l_+(int)(h_),:]=(200,200,200)
            if vel_ego[0][0][2]>=0 :
                h_ = l_*(vel_ego[0][0][2]/self.spd)
                h_ = min(l_,h_)
                observation[:12,l_*2:l_*2+(int)(h_),:]=(200,200,200)
            elif vel_ego[0][0][2]<0 :
                h_ = -l_*(vel_ego[0][0][2]/self.spd)
                h_ = min(l_,h_)
                observation[:12,l_*3:l_*3+(int)(h_),:]=(200,200,200)

            self.camera_sensor_image = np.copy(observation)

        else:
            observation = np.zeros((HEIGHT,WIDTH, 3), np.uint8)

        '''
        get reward
        '''
        reward = self.calculate_reward(self.envs[i], self.actor_handles[i], self.food_handles[i])
        done = self.done_determinant(self.envs[i], self.actor_handles[i], self.food_handles[i],reward,self.step_count)
        info = {}

        if self.show_action:
            self.show_image()
        
        if not self.headless:
            # update the viewer
            self.gym.draw_viewer(self.viewer, self.sim, True)

        # end_time = time.time()
        # if (end_time - self.start_time) > 1:
        #     print("FPS: %.2f" % ((self.step_count - self.last_frame_cnt)*self.num_matrix_envs))
        #     self.last_frame_cnt = self.step_count
        #     self.start_time = time.time()
        
        return np.copy(observation), reward, done, info

    def reset(self):

        self.step_count = 0
        self.facing = 0
        self.start_time = time.time()

        i = 0

        done_at_begining = False #self.done_determinant(self.envs[i], self.actor_handles[i], self.food_handles[i])

        while not done_at_begining:
            
            self.gym.set_env_rigid_body_states(self.envs[i], self.sim_status[i], gymapi.STATE_ALL)

            # reset random position and orientation for actor
            x = random.uniform(0,16)
            y = random.uniform(0,16)
            facing_ = 0.25 * (float)(random.randint(0,8))
            r_ = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing_*math.pi)
            body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
            body_states['pose']['p'].fill((x,y,self.initial_height))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_POS)
            attractor_pose = gymapi.Transform()
            attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
            attractor_pose.r = r_
            self.gym.set_attractor_target(self.envs[i], self.attractor_handles[i], attractor_pose)
            # print("reset actor at (%d, %d)" % (x, y))
            # reset random position for food
            x = random.uniform(0,16)
            y = random.uniform(0,16)
            body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.food_handles[i], gymapi.STATE_ALL)
            body_states['pose']['p'].fill((x,y,self.initial_height))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.food_handles[i], body_states, gymapi.STATE_POS)

            # reset actor vel=0
            body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
            body_states['vel']['linear'].fill((0,0,0))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)

            # step the physics
            self.gym.simulate(self.sim)
            self.gym.fetch_results(self.sim, True)
            self.gym.step_graphics(self.sim)
            self.gym.render_all_camera_sensors(self.sim)

            rgb_image = self.gym.get_camera_image(self.sim, self.envs[i], self.camera_handles[i], gymapi.IMAGE_COLOR)
            # create an opencv image of size (HEIGHT,WIDTH , 4)
            observation = np.zeros((HEIGHT,WIDTH, 4), np.uint8)
            observation[:, :, :4] = rgb_image.reshape((HEIGHT,WIDTH, 4))
            # rotate 90 degrees counter-clockwise
            observation = np.rot90(observation, -1)
            # transform the image from RGBA to RGB
            observation = observation[:, :, :3]

            '''
            simulate until actor and food lands on the terrain
            '''
            for t in range(40):
                observation, reward, done, info = self.step(action = -3, physics_only = True)

            done_at_begining = not self.done_determinant(self.envs[i], self.actor_handles[i], self.food_handles[i])
        
        observation, reward, done, info = self.step(action = -3, physics_only = False)
        self.camera_sensor_image = np.copy(observation)

        return np.copy(observation)

    def render(self, mode='human'):
        
        i = 0
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states['pose']['p']
        # print("actor position: (%f, %f, %f)" % (actor_p[0][0], actor_p[0][1], actor_p[0][2]))
        cnnc_x = ((actor_p[0][0]+1)/(9*2))*20*9
        cnnc_y = ((actor_p[0][1]+1)/(9*2))*20*9

        # print(self.connectivity_map[int(cnnc_y)][int(cnnc_x)])

        if not self.headless:

            if self.gym.query_viewer_has_closed(self.viewer):
                return -2

            # for i in range(self.num_matrix_envs):
            body_states = self.gym.get_camera_transform(self.sim, self.envs[i], self.camera_handles[i])
            draw_camera(self.gym, self.viewer,self.envs[i],body_states)
            # self.gym.draw_env_rigid_contacts(self.viewer, self.envs[i], gymapi.Vec3(1,0,0), 2, True)
            # update the viewer
            self.gym.draw_viewer(self.viewer, self.sim, True)

            self.camera_sensor_image = cv2.cvtColor(self.camera_sensor_image, cv2.COLOR_RGB2BGR)
            cv2.imshow('rgb', self.camera_sensor_image)

            # draw dot at [cnnc_x, cnnc_y]
            cnnc_vis = np.copy(self.connectivity_map_visual)
            cv2.circle(cnnc_vis, (int(cnnc_x), int(cnnc_y)), 5, (0, 0, 255), -1)
            cv2.imshow('conn',cnnc_vis)
            self.interaction = cv2.waitKey(1)
            
        return self.interaction

    def show_image(self):
        
        i = 0
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states['pose']['p']
        # print("actor position: (%f, %f, %f)" % (actor_p[0][0], actor_p[0][1], actor_p[0][2]))
        cnnc_x = ((actor_p[0][0]+1)/(9*2))*20*9
        cnnc_y = ((actor_p[0][1]+1)/(9*2))*20*9

        self.camera_sensor_image = cv2.cvtColor(self.camera_sensor_image, cv2.COLOR_RGB2BGR)
        cv2.imshow('rgb', self.camera_sensor_image)
        cv2.waitKey(1)

    def close (self):
        
        if not self.headless:
            self.gym.destroy_viewer(self.viewer)

        self.gym.destroy_sim(self.sim)
